from pro_authformer.src.main.servlet_main.config import config_dict
from pro_authformer.src.main.servlet_component.e_loss import Loss
from pro_authformer.src.main.servlet_component.e_optimizer import Optimizer
from pro_authformer.src.main.servlet_component.e_scheduler import Scheduler

class Trainer:
    def __init__(self):
        # 初始化损失函数
        self.loss = Loss()
        # 初始化优化器
        self.optimizer = Optimizer()
        # 初始化学习率调度器
        self.scheduler = Scheduler()
    def train_setup(self):
        loss = self.loss.get_loss(config_dict["loss_type"])
        optimizer = self.optimizer.get_optimizer(config_dict["optimizer_type"])
        scheduler = self.scheduler.get_scheduler(config_dict["scheduler_type"])
        
        return loss, optimizer, scheduler
    
    
    def fit_one_epoch(self, model, train_dataset, val_dataset, loss, optimizer, scheduler):
        pass